package cn.wei.batch;

import org.apache.rocketmq.common.message.Message;

import java.util.Iterator;
import java.util.List;
import java.util.Map;

/**
 * 消息列表分割器：其只会处理每条消息的大小不超过4M的情况
 * 若存在某条消息，其本身大小大于4M，则消息列表分割器无法处理
 * 其直接将这条消息构成一个子列表返回，不会再去分割
 */
public class MessageListSplitter implements Iterator<List<Message>> {

    // 指定消费极限值为4M
    private final int SIZE_LIMIT = 1024 * 4 * 1024;

    // 存放所有要发送的消息
    private final List<Message> messages;

    // 批量发送消息的小集合起始索引
    private int currIndex;

    public MessageListSplitter(List<Message> messages) {
        this.messages = messages;
    }

    @Override
    public boolean hasNext() {
        // 判断当前开始的消息索引要小于消息总数
        return currIndex < messages.size();
    }

    @Override
    public List<Message> next() {
        int nextIndex = currIndex;
        // 记录这一次要发送消息的列表大小
        int totalSize = 0;
        for (; nextIndex < messages.size(); nextIndex++) {
            // 获取当前遍历消息
            Message message = messages.get(nextIndex);
            // 统计当前遍历的message的大小
            int tmpSize = message.getTopic().length() + message.getBody().length;
            Map<String, String> properties = message.getProperties();
            for (Map.Entry<String, String> entry : properties.entrySet()) {
                tmpSize += entry.getKey().length() + entry.getValue().length();
            }
            tmpSize += 20;

            // 判断当前消息本身是否大于4M
            if (tmpSize > SIZE_LIMIT) {
                if (nextIndex - currIndex == 0) {
                    nextIndex++;
                }
                break;
            }
            if (tmpSize + totalSize > SIZE_LIMIT) {
                break;
            } else {
                totalSize += tmpSize;
            }
        }
        // 获取当前message列表的子集合[currIndex,nextIndex]
        List<Message> messageList = this.messages.subList(currIndex, nextIndex);
        // 下次开始遍历时的索引
        currIndex = nextIndex;
        return messageList;
    }
}
